import json

from .utils.logger import get_logger

from fastapi import APIRouter, HTTPException
from .api.payload import BacktrackingTask, DeleteTask
from ray.job_submission import JobSubmissionClient

backtracking_router = APIRouter()
logger = get_logger()


def filter_backtracking_results(backtracking_results, delete_intervals, target_noid, changed_info, deleted_set,
                                query_info):
    for noid, value in backtracking_results.items():
        if noid != target_noid:
            continue
        outputs = value["result"]
        sims = value["sims"]

        # 标记需要删除的索引
        to_delete_indices = []

        for i, (output, sim) in enumerate(zip(outputs, sims)):
            start, end = output
            if [start, end] in delete_intervals:
                to_delete_indices.append(i)

        if to_delete_indices:
            # 检查是否有源节点被删除
            if any(sims[i] >= 1 for i in to_delete_indices):
                # 如果有源节点被删除，删除整个结果
                del backtracking_results[noid]
                query_info = [q for q in query_info if q["noid"] != noid]
                changed_info["deleted_set"].append(noid)
                deleted_set.append(noid)
                break

            # 否则，只删除从节点
            value["result"] = [output for i, output in enumerate(outputs) if i not in to_delete_indices]
            value["sims"] = [sim for i, sim in enumerate(sims) if i not in to_delete_indices]
            changed_info["updated_set"].append(noid)


@backtracking_router.get("/")
def root():
    return "Thanks for using backtraking offline service!"


@backtracking_router.post('/submit/job')
def submit_task(request: BacktrackingTask):
    try:
        # data = request.model_dump_json().encode('utf-8')
        data = request.model_dump_json()
        client = JobSubmissionClient("http://localhost:8265")
        job_id = client.submit_job(
            # Entrypoint shell command to execute
            entrypoint=f"cd backtracking && python backtracking.py -i \'{data}\' -c configs/config.yml",
            # entrypoint=f"echo hello",
            # submission_id= task_id,
            # Path to the local directory that contains the script.py file
            runtime_env={"working_dir": "/work/job"
                         },

        )
        logger.info({"job_id": job_id, "dashboard": "http://localhost:8265"})
        return {"job_id": job_id, "dashboard": "http://localhost:8265"}
    except Exception as e:
        logger.error(f"auto match task occurs: {e}")
        return {"msg": str(e)}


@backtracking_router.post('/delete')
def delete_data(request: DeleteTask):
    task_id = request.task_id
    ranges = request.delete_ranges
    target_noid = request.noid
    try:
        with open(request.meta_path, 'r') as f:
            meta_data = json.load(f)
            print(meta_data)
    except FileNotFoundError as e:
        logger.error(f"could not find meta file for {task_id}")
        raise HTTPException(status_code=404, detail=f"could not find meta file for {task_id}")

    try:
        if "versions" not in meta_data.keys():
            raise HTTPException(status_code=404, detail=f"could not find versions in meta file for {task_id}")
        query_info = meta_data["query_info"]
        versions = meta_data["versions"]
        versions["changed_info"] = {}
        changed_info = versions["changed_info"]
        changed_info["deleted_set"] = []
        changed_info["updated_set"] = []

        deleted_set = versions.get("deleted_set", [])
        backtracking_results = versions.get("backtracking_results", [])
        print(backtracking_results)
        filter_backtracking_results(backtracking_results, ranges, target_noid, changed_info,
                                    deleted_set, query_info)
        meta_data['update_status'] = 'update_version'
        # backtracking_path = os.path.splitext(ref_meta["path"])[0] + "-backtracking.json"
        with open(request.meta_path, "w", encoding="utf-8") as f:
            json_data = json.dumps(meta_data, indent=4, ensure_ascii=True)
            f.write(json_data)

        return {"status": "success", "data": request.meta_path}
    except HTTPException as e:
        raise e
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
